import pandas as pd
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import numpy as np
from dataset.utils import FRegression
class Data(Dataset):
    def __init__(self, season, sheet_name, step_back=False, k=3):
        season = season
        sheet_name = sheet_name
        df = pd.read_excel('dataset/data/%s.xlsx' % season, sheet_name=sheet_name)
        df = df.dropna()

        self.columns = df.loc[:, ['空气温度', '十分风速', '湿度', '气压']].columns
        self.date = df.date
        self.inputs = df.loc[:, ['空气温度', '十分风速', '湿度', '气压']].values.astype("float32")
        self.targets = df.loc[:, 'TA_CU'].values.astype("float32")

        self.support = [True]*4
        if step_back:
            self.inputs, self.support = FRegression(self.inputs, self.targets)
        self.columns = self.columns[self.support].tolist()

    def __getitem__(self, item):

        return self.date[item], self.inputs[item], self.targets[item]

    def __len__(self):
        return len(self.date)

    def get_support(self):
        return self.support

    def get_columns(self):

        return self.columns, ['TA_CU']

class CreateDataset:
    def __init__(self, args):
        self.batch_size = args.batch_size
        self.step_back = args.step_back
        self.k = args.n

    def getDataLoader(self, season, sheet_name):
        self.data = Data(season, sheet_name, self.step_back, self.k)
        dataLoader = DataLoader(self.data, shuffle=False, batch_size=self.batch_size)

        return dataLoader

    def getcolumns(self):

        return self.data.get_columns()

    def get_support(self):

        return self.data.support
